package problem208_Implement_Trie;

class TrieNode {
	char val;
	TrieNode[] sons;
	boolean isEnd;

	public TrieNode() {
		sons = new TrieNode[26];
		isEnd = false;
	}

	public TrieNode(char c) {
		sons = new TrieNode[26];
		val = c;
		isEnd = false;
	}
}

public class Trie {
	private TrieNode root;

	public Trie() {
		root = new TrieNode();
	}

	// Inserts a word into the trie.
	public void insert(String word) {
		char[] arr = word.toCharArray();
		TrieNode node = root;
		for (char c : arr) {
			int pos = c - 'a';
			if (node.sons[pos] == null) {
				node.sons[pos] = new TrieNode(c);
			}
			node = node.sons[pos];
		}
		node.isEnd = true;
	}

	// Returns if the word is in the trie.
	public boolean search(String word) {
		char[] arr = word.toCharArray();
		TrieNode node = root;
		for (char c : arr) {
			int pos = c - 'a';
			if (node.sons[pos] == null) {
				return false;
			} else {
				node = node.sons[pos];
			}
		}
		return node.isEnd;
	}

	// Returns if there is any word in the trie
	// that starts with the given prefix.
	public boolean startsWith(String prefix) {
		char[] arr = prefix.toCharArray();
		TrieNode node = root;
		for (char c : arr) {
			int pos = c - 'a';
			if (node.sons[pos] == null) {
				return false;
			}
			node=node.sons[pos];
		}
		return true;
	}
	
	public static void main(String[] args){
		Trie trie=new Trie();
		trie.insert("a");
		trie.insert("ab");
		System.out.println(trie.search("a"));
		System.out.println(trie.search("ab"));
		System.out.println(trie.startsWith("ab"));
	}
}